# import the libraries
import os
from pathlib import Path

import cv2
import numpy as np

from conf import myConfig as config

# define scales
# scales=[1,0.9,0.8,0.7,0.6]
scales = [1, 0.9, 0.8, 0.7]
count = 0

# create the imageArrays
p = Path(config.genDataPath)
listPaths = list(p.glob('./*.png'))
imgArray = []
for path in listPaths:
    imgArray.append(cv2.imread(str(path), 0))
print('lenImages', len(imgArray))

print(imgArray[0].shape)

# calculate the number of patches
for i in range(len(imgArray)):
    img = imgArray[i]
    for s in range(len(scales)):
        newsize = (int(img.shape[0] * scales[s]), int(img.shape[1] * scales[s]))
        img_s = cv2.resize(img, newsize, interpolation=cv2.INTER_CUBIC)
        im_h, im_w = img_s.shape
        for x in range(0 + config.step, (im_h - config.pat_size), config.stride):
            for y in range(0 + config.step, (im_w - config.pat_size), config.stride):
                count += 1

origin_patch_num = count
if origin_patch_num % config.batch_size != 0:
    numPatches = (origin_patch_num / config.batch_size + 1) * config.batch_size
else:
    numPatches = origin_patch_num
print('total patches=%d, batch_size=%d, total_batches=%d' %
      (numPatches, config.batch_size, numPatches / config.batch_size))

# numpy array to contain patches for training
inputs = np.zeros((int(numPatches), int(config.pat_size), int(config.pat_size), 1), dtype=np.uint8)

# generate patches
count = 0
for i in range(len(imgArray)):
    img = imgArray[i]
    for s in range(len(scales)):
        newsize = (int(img.shape[0] * scales[s]), int(img.shape[1] * scales[s]))
        img_s = cv2.resize(img, newsize, interpolation=cv2.INTER_CUBIC)
        img_s = np.reshape(np.array(img_s, dtype="uint8"),
                           (img_s.shape[0], img_s.shape[1], 1))
        im_h, im_w, _ = img_s.shape
        for x in range(0 + config.step, im_h - config.pat_size, config.stride):
            for y in range(0 + config.step, im_w - config.pat_size, config.stride):
                inputs[count, :, :, :] = img_s[x:x + config.pat_size, y:y + config.pat_size, :]
                count += 1

# pad the batch
if count < numPatches:
    to_pad = int(numPatches - count)
    inputs[-to_pad:, :, :, :] = inputs[:to_pad, :, :, :]

if not os.path.exists(config.save_dir):
    os.mkdir(config.save_dir)
np.save(os.path.join(config.save_dir, "img_clean_pats"), inputs)
